package com.anttribe.components.db.metadata.supplier;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import javax.sql.DataSource;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import com.anttribe.components.db.metadata.Dialect;
import com.anttribe.components.db.metadata.datasource.DataSourceProxy;
import com.anttribe.components.db.metadata.entity.Database;
import com.anttribe.components.db.metadata.entity.DbColumn;
import com.anttribe.components.db.metadata.entity.DbIndex;
import com.anttribe.components.db.metadata.entity.DbIndexType;
import com.anttribe.components.db.metadata.entity.DbObjectType;
import com.anttribe.components.db.metadata.entity.DbSchema;
import com.anttribe.components.db.metadata.entity.DbTable;
import com.anttribe.components.db.metadata.entity.JdbcType;
import com.anttribe.components.db.metadata.entity.PrimaryKey;
import com.anttribe.components.db.metadata.exception.DbException;
import com.anttribe.components.db.metadata.utils.DbUtils;
import com.google.common.collect.Lists;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

/**
 * @author zhaoyong
 * @date 2021/10/21 0021
 */

@Slf4j
public abstract class AbstractDatabaseSupplier implements DatabaseSupplier {

    /**
     * 表名通配符
     */
    private static final String DEFAULT_TABLENAME_WILDCARD = "%";

    @Getter
    private final DataSource dataSource;

    /**
     * 数据库连接
     */
    private Connection connection;

    public AbstractDatabaseSupplier(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    /**
     * 获取数据库连接
     *
     * @return Connection
     */
    protected Connection getConnection() {
        try {
            if (null != this.connection && !this.connection.isClosed()) {
                return this.connection;
            }

            this.connection = this.dataSource.getConnection();
            return connection;
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

    /**
     * 获取数据库元数据信息
     *
     * @return DatabaseMetaData
     */
    private DatabaseMetaData getDatabaseMetaData() {
        return this.getDatabaseMetaData(this.getConnection());
    }

    /**
     * 获取当前处理数据库方言
     *
     * @return Dialect
     */
    @Override
    public Dialect getDialect() {
        DataSourceProxy dataSourceProxy = (DataSourceProxy)this.getDataSource();
        try {
            return dataSourceProxy.getDialect();
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

    /**
     * 获取数据库元数据信息
     *
     * @param conn
     *            Connection
     * @return DatabaseMetaData
     */
    private DatabaseMetaData getDatabaseMetaData(Connection conn) {
        if (null == conn) {
            conn = this.getConnection();
        }

        try {
            return conn.getMetaData();
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

    @Override
    public Database getDatabase() {
        try {
            // 获取数据库元数据信息
            DatabaseMetaData databaseMetaData = this.getDatabaseMetaData();

            Database database = new Database();
            database.setProductName(databaseMetaData.getDatabaseProductName());
            database.setProductVersion(databaseMetaData.getDatabaseProductVersion());
            return database;
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

    @Override
    public DbSchema getDbSchema() {
        List<DbSchema> dbSchemas = this.getAvailableDbSchemas();
        if (CollectionUtils.isEmpty(dbSchemas)) {
            return null;
        }

        try {
            // 从当前连接中获取DbSchema
            DbSchema curDbSchema = new DbSchema();
            curDbSchema.setCatalog(this.getConnection().getCatalog());
            curDbSchema.setSchema(this.getConnection().getSchema());
            Optional<DbSchema> curDbSchemaOptional = dbSchemas.stream()
                .filter(dbSchema -> StringUtils.equals(dbSchema.getFullName(), curDbSchema.getFullName())).findAny();
            return curDbSchemaOptional.orElseGet(null);
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

    /**
     * 获取当前数据库连接可用的DbSchema
     *
     * @return List<DbSchema>
     */
    private List<DbSchema> getAvailableDbSchemas() {
        List<DbSchema> dbSchemas = new ArrayList<DbSchema>();

        DatabaseMetaData databaseMetaData = this.getDatabaseMetaData();
        List<String> catalogs = this.getAvailableDbCatalogs();
        if (CollectionUtils.isEmpty(catalogs)) {
            try (ResultSet schemasRs = databaseMetaData.getSchemas()) {
                while (schemasRs.next()) {
                    DbSchema schema = new DbSchema();
                    schema.setCatalog(schemasRs.getString("TABLE_CATALOG"));
                    schema.setSchema(schemasRs.getString("TABLE_SCHEM"));
                    dbSchemas.add(schema);
                }
            } catch (SQLException e) {
                throw new DbException(e);
            }
        } else {
            try {
                for (String catalog : catalogs) {
                    List<DbSchema> tempSchemas = new ArrayList<DbSchema>();
                    ResultSet schemasRs = databaseMetaData.getSchemas(catalog, null);
                    while (schemasRs.next()) {
                        DbSchema schema = new DbSchema();
                        schema.setCatalog(schemasRs.getString("TABLE_CATALOG"));
                        schema.setSchema(schemasRs.getString("TABLE_SCHEM"));
                        tempSchemas.add(schema);
                    }
                    DbUtils.closeQuietly(schemasRs);

                    if (CollectionUtils.isEmpty(tempSchemas)) {
                        DbSchema schema = new DbSchema();
                        schema.setCatalog(catalog);
                        dbSchemas.add(schema);
                    } else {
                        dbSchemas.addAll(tempSchemas);
                    }
                }
            } catch (SQLException e) {
                throw new DbException(e);
            }
        }
        return dbSchemas;
    }

    /**
     * 获取当前数据库连接可用的DbCatalogs
     *
     * @return List<String>
     */
    private List<String> getAvailableDbCatalogs() {
        DatabaseMetaData databaseMetaData = this.getDatabaseMetaData();
        // 设置 catalogs
        List<String> catalogs = new ArrayList<String>();
        try (ResultSet rs = databaseMetaData.getCatalogs()) {
            while (rs.next()) {
                catalogs.add(rs.getString("TABLE_CAT"));
            }
            return catalogs;
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

    @Override
    public List<DbTable> getTables() {
        return getTables(DEFAULT_TABLENAME_WILDCARD);
    }

    @Override
    public List<DbTable> getTables(String tableNamePattern) {
        return this.getTablesInner(tableNamePattern);
    }

    @Override
    public DbTable getTable(String tableName) {
        if (StringUtils.isEmpty(tableName)) {
            return null;
        }

        List<DbTable> dbTables = this.getTablesInner(tableName);
        if (CollectionUtils.isEmpty(dbTables)) {
            return null;
        }

        for (DbTable dbTable : dbTables) {
            if (StringUtils.equalsIgnoreCase(tableName, dbTable.getTableName())) {
                return dbTable;
            }
        }

        return null;
    }

    /**
     * 获取数据库表
     *
     * @param tableNamePattern
     *            表名
     * @return List<DbTable>
     */
    private List<DbTable> getTablesInner(String tableNamePattern) {
        DbSchema dbSchema = this.getDbSchema();
        if (null == dbSchema) {
            return null;
        }

        List<DbTable> dbTables = new ArrayList<DbTable>();
        DatabaseMetaData databaseMetaData = this.getDatabaseMetaData();
        try (ResultSet rs = databaseMetaData.getTables(dbSchema.getCatalog(), dbSchema.getSchema(), tableNamePattern,
            new String[] {DbObjectType.TABLE.name()})) {
            while (rs.next()) {
                DbTable dbTable = new DbTable();
                dbTable.setDbSchema(dbSchema);
                dbTable.setTableName(rs.getString("TABLE_NAME"));
                dbTable.setComment(rs.getString("REMARKS"));
                dbTable.setType(DbObjectType.TABLE);

                // 数据列
                List<DbColumn> dbColumns = this.getTableColumns(dbTable);
                dbTable.setColumns(dbColumns);
                // 主键
                PrimaryKey primaryKey = this.getPrimaryKey(dbTable);
                dbTable.setPrimaryKey(primaryKey);
                // 将主键设置到列信息中
                List<String> primaryKeyColumns = primaryKey.getColumns();
                if (!CollectionUtils.isEmpty(primaryKeyColumns) && !CollectionUtils.isEmpty(dbColumns)) {
                    dbColumns.forEach(dbColumn -> {
                        if (primaryKeyColumns.contains(dbColumn.getColumnName())) {
                            dbColumn.setPrimaryKey(Boolean.TRUE);
                        }
                    });
                }
                // 索引
                dbTable.setIndexes(this.getTableIndexes(dbTable));

                dbTables.add(dbTable);
            }

            return dbTables;
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

    /**
     * 查询数据库表下的字段
     *
     * @param dbTable
     *            数据库表信息
     * @return List<DbColumn>
     */
    private List<DbColumn> getTableColumns(DbTable dbTable) {
        DbSchema dbSchema = dbTable.getDbSchema();
        if (null == dbSchema) {
            return null;
        }

        List<DbColumn> dbColumns = new ArrayList<DbColumn>();
        DatabaseMetaData databaseMetaData = this.getDatabaseMetaData();
        try (ResultSet columnsRs =
            databaseMetaData.getColumns(dbSchema.getCatalog(), dbSchema.getSchema(), dbTable.getTableName(), null)) {
            while (columnsRs.next()) {
                DbColumn dbColumn = new DbColumn();
                dbColumn.setColumnName(columnsRs.getString("COLUMN_NAME"));
                dbColumn.setComment(columnsRs.getString("REMARKS"));
                dbColumn.setJdbcType(JdbcType.valueOf(columnsRs.getInt("DATA_TYPE")));
                dbColumn.setSize(columnsRs.getInt("COLUMN_SIZE"));
                dbColumn.setPrecision(dbColumn.getSize());
                dbColumn.setDigits(columnsRs.getInt("DECIMAL_DIGITS"));
                dbColumn.setNullable(columnsRs.getInt("NULLABLE") == ResultSetMetaData.columnNullable);
                dbColumn.setDefaultVal(columnsRs.getString("COLUMN_DEF"));
                dbColumn.setOrdinal(columnsRs.getInt("ORDINAL_POSITION"));
                dbColumns.add(dbColumn);
            }
            dbColumns =
                dbColumns.stream().sorted(Comparator.comparing(DbColumn::getOrdinal)).collect(Collectors.toList());
            return dbColumns;
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

    /**
     * 获取主键信息
     *
     * @param dbTable
     *            DbTable
     * @return PrimaryKey
     */
    private PrimaryKey getPrimaryKey(DbTable dbTable) {
        DbSchema dbSchema = dbTable.getDbSchema();
        if (null == dbSchema) {
            return null;
        }

        PrimaryKey primaryKey;
        DatabaseMetaData databaseMetaData = this.getDatabaseMetaData();
        try (ResultSet primaryKeyRs =
            databaseMetaData.getPrimaryKeys(dbSchema.getCatalog(), dbSchema.getSchema(), dbTable.getTableName())) {
            primaryKey = new PrimaryKey();
            while (primaryKeyRs.next()) {
                primaryKey.setName(primaryKeyRs.getString("PK_NAME"));
                primaryKey.addColumn(primaryKeyRs.getString("COLUMN_NAME"), primaryKeyRs.getShort("KEY_SEQ"));
            }
            return primaryKey;
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

    /**
     * 获取数据库索引信息
     *
     * @param dbTable
     *            数据库表
     * @return List<DbIndex>
     */
    private List<DbIndex> getTableIndexes(DbTable dbTable) {
        DbSchema dbSchema = dbTable.getDbSchema();
        if (null == dbSchema) {
            return null;
        }

        DatabaseMetaData databaseMetaData = this.getDatabaseMetaData();
        try (ResultSet indexRs = databaseMetaData.getIndexInfo(dbSchema.getCatalog(), dbSchema.getSchema(),
            dbTable.getTableName(), Boolean.FALSE, Boolean.TRUE)) {
            Map<String, DbIndex> dbIndexMap = new HashMap();
            while (indexRs.next()) {
                String indexName = indexRs.getString("INDEX_NAME");
                DbIndex dbIndex = dbIndexMap.getOrDefault(indexName, new DbIndex());
                dbIndex.setIndexName(indexName);
                dbIndex.setUnique(indexRs.getBoolean("NON_UNIQUE"));
                dbIndex.setIndexType(DbIndexType.valueOfType(indexRs.getShort("TYPE")));
                String columnName = indexRs.getString("COLUMN_NAME");
                if (StringUtils.isNotEmpty(columnName)) {
                    List<String> indexColumns = Optional.ofNullable(dbIndex.getColumns()).orElse(Lists.newArrayList());
                    indexColumns.add(columnName);
                    dbIndex.setColumns(indexColumns);
                }
                dbIndexMap.put(indexName, dbIndex);
            }

            List<DbIndex> dbIndexes = new ArrayList<>();
            dbIndexes.addAll(dbIndexMap.values());
            return dbIndexes;
        } catch (SQLException e) {
            throw new DbException(e);
        }
    }

}
